# 挖掘子图3
import csv
import json

# 读取Node.csv
node_file = '../Data/Node.csv'
node_dict = {}
with open(node_file, 'r', encoding='utf-8') as f:
    reader = csv.reader(f)
    next(reader)
    for row in reader:
        node_id, node_name, node_type, node_industry = row
        node_dict[node_id] = {'node_id': node_id,'name': node_name, 'type': node_type, 'industry': node_industry}

# 读取Link.csv
link_file = '../Data/Link.csv'
link_dict = {}
with open(link_file, 'r', encoding='utf-8') as f:
    reader = csv.reader(f)
    next(reader)
    for row in reader:
        relation, source, target = row
        if source not in link_dict:
            link_dict[source] = {}
        if relation not in link_dict[source]:
            link_dict[source][relation] = []
        link_dict[source][relation].append(target)

# 定义函数
def explore(node_ids):
    explored_nodes = set()
    explored_links = set()
    nodes_to_explore = set(node_ids)
    for i in range(3):
        new_nodes_to_explore = set()
        for node_id in nodes_to_explore:
            if node_id in explored_nodes:
                continue
            explored_nodes.add(node_id)
            for relation, targets in link_dict.get(node_id, {}).items():
                if relation in ['r_cert', 'r_subdomain', 'r_request_jump', 'r_dns_a'] and i < 3:
                    new_nodes_to_explore.update(targets)
                elif relation in ['r_whois_name', 'r_whois_email', 'r_whois_phone'] and i < 2:
                    new_nodes_to_explore.update(targets)
                elif relation in ['r_cert_chain', 'r_cname'] and i < 1:
                    new_nodes_to_explore.update(targets)
                elif relation in ['r_asn', 'r_cidr'] and i < 1:
                    new_nodes_to_explore.update(targets)
                for target in targets:
                    explored_links.add((node_id, relation, target))
        nodes_to_explore = new_nodes_to_explore
        if not nodes_to_explore:
            break
    return explored_nodes, explored_links

# 处理多个初始节点
start_nodes = ['Domain_24acfd52f9ceb424d4a2643a832638ce1673b8689fa952d9010dd44949e6b1d9',
               'Domain_9c72287c3f9bb38cb0186acf37b7054442b75ac32324dfd245aed46a03026de1',
               'Domain_717aa5778731a1f4d6f0218dd3a27b114c839213b4af781427ac1e22dc9a7dea',
               'Domain_8748687a61811032f0ed1dcdb57e01efef9983a6d9c236b82997b07477e66177',
               'Whois_Phone_f4a84443fb72da27731660695dd00877e8ce25b264ec418504fface62cdcbbd7']
final_nodes = set()
final_links = set()
for start_node in start_nodes:
    nodes, links = explore([start_node])
    final_nodes.update(nodes)
    final_links.update(links)

# 将结果输出到json文件
output_file = 'subgraph3.json'
with open(output_file, 'w', encoding='utf-8') as f:
    result = {
        'nodes': {node_id: node_dict[node_id] for node_id in final_nodes},
        'links': [{'source': source, 'relation': relation, 'target': target} for source, relation, target in final_links]
    }
    json.dump(result, f, ensure_ascii=False, indent=4)
